最近一次想要批量运行任务前试着运行了一个个例，发现占用的内存意外的比较多，所以一口气把所有提交到服务器上显然不太妥当。由于之前在服务器上用 slurm 仅限于 srun 而从来没有尝试过 sbatch，所以就对这方面的知识进行了一番恶补，正好也能补充十二月份的博客内容。

如何确定内存是否充足

对于当前节点，使用 top 和 htop 就能得到相关的信息，但对于其他节点，如果没有登录权限，那么这种实时监控的手段就没法派上用场了。

这时候，可以使用 free -h 命令将节点在该时间点的内存使用情况信息打印出来：

1
2
3
4
$ srun -w nodexx free -h
              total        used        free      shared  buff/cache   available
Mem:           487G        144G        2.7G        162M        340G        340G
Swap:           49G        685M         49G

重点关注 available 部分，可以看到当前该节点还是剩余了非常多可用内存的。不过并不是内存越多可以运行的命令就越多，这也跟剩余的空闲 CPU 数有关，关于剩余的 CPU 数可以直接使用 scontrol show node nodexx 来查看：

1
2
$ scontrol show node nodexx | grep 'CPUTot'
   CPUAlloc=45 CPUTot=52 CPULoad=75.96

可以看到已经分配的 CPU 有多少个（CPUAlloc）。

如何确定任务要多少内存

可以先只投递单个任务，若该任务投递到当前终端所在节点，则直接使用 top -u username 就能进行查看。

如果该任务投递到其他节点上，那么可以使用 ps 命令查看相关情况：

1
2
3
$ srun -w nodexx ps -u username -o pid,vsz,rss,comm
   PID    VSZ   RSS COMMAND
  6382 34042040 9706368 xxxx

这里我们可以着重关注 RSS（常驻内存集大小，即进程实际使用的物理内存大小），该例子中这个命令占用了约 9.7 GB（9706368 KB） 的内存，以此为准，若我们剩余了 100 GB 的内存，那就还能支撑我们运行大约十个这样的命令（理想情况）。但实际情况肯定没那么简单，所以最好留出一些余量以避免性能下降或内存交换。

使用 sbatch 提交作业

确定了自己大约还能投多少任务后，就可以针对其编写相关的 shell 脚本了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#!/bin/bash

#SBATCH --job-name=xxx
#SBATCH --output=xxx_%A_%a.out
#SBATCH --error=xxx_%A_%a.err
#SBATCH --partition=xxx
#SBATCH --ntasks=1
#SBATCH --array=1-10%5

FILES=(*txt)
INDEX=$(($SLURM_ARRAY_TASK_ID - 1))
FILE=${FILES[$INDEX]}

echo $FILE

以上述例子为例，以下几个标头较为重要：

  • #SBATCH --output=xxx_%A_%a.out：这里的 %A 和 %a 分别表示作业数组的 ID 和数组中任务的索引，添加后不同作业的 out 和 err 会分开产生而不会混在一起。

  • #SBATCH --partition=xxx：该处的填写可根据 sinfo 的信息进行，例如：

1
2
3
$ sinfo
PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
all*         up   infinite      3    mix nodexx,nodexx,nodexx

此时可填写 all，slurm 会根据该 partition 中 node 的使用情况自动分配任务。如果没有 partition 可以换成 #SBATCH --nodelist=xxx 填写想要投递的节点名称。

  • #SBATCH --ntasks=1：这行声明了每个作业使用的 CPU 数量，此处为每个作业使用 1 个。需要注意的是，提交的作业数组任务其本质上为多个任务的组合，因此添加的各类资源限制都是对每个任务单独生效的。因此此处同时运行 5 个作业（%5）则需要 5 个 CPU。同样地，也可以使用 --time 和 --mem 等选项对单个作业的运行进行限制。
  • #SBATCH --array=1-10%5：这行定义了作业数组的范围（1-10，共十个）和同时运行的作业数量（%5），这里还有其他的设计技巧，例如：
    • 1-100:10 表示从 1 到 100，每隔 10 个一个作业。
    • 1,2,5-7 表示作业 1, 2, 5, 6, 和 7。
    • slurm 为每个作业数组提供了特定的环境变量，如示例脚本中的 $SLURM_ARRAY_TASK_ID

以上脚本会按照以下逻辑进行：

  1. FILE 变量储存当前文件夹下所有 txt 结尾的文件。
  2. 通过 $SLURM_ARRAY_TASK_ID 定义索引变量 INDEX 并调取出 FILE 中对应索引的文件。
  3. 打印出对应文件的文件名（echo），结果将出现在 out 文件中。

运行该脚本的方法（假设该脚本名字为 test.sh）：

1
2
$ sbatch test.sh
$ sbatch --mail-user=xxx@xxx --mail-type=BEGIN,END test.sh # 开始和结束后发邮件通知

上述代码中，发邮件通知也可以通过在脚本中添加以下代码完成：

1
2
#SBATCH --mail-user=xxx@xxx
#SBATCH --mail-type=BEGIN,END

投递任务后，如果想要关注相关任务的状态，可以使用 squeue -u username 进行。

当然，如果需要提交的作业不能仅仅通过索引来区分，而是需要执行完全不同的命令时，使用作业数组可能不是最合适的选择。这时候对每个任务单独编写脚本并提交更能满足实际需求。